# -*- coding: utf-8 -*-
# coding=utf-8
# 测试
import sys
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '3'
# 获取父目录的绝对路径
parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
# 将父目录添加到sys.path
sys.path.append(parent_dir)
from modules.literal_recognition import LiteralRecognition
from utils.chatglm_utils import ChatGLMUtils
model = ChatGLMUtils()
model.load_model(local_loading=True,
                 model_path=f"models/all-chatglm-6b-lora.pt")
if __name__ == "__main__":
    literal_reg = LiteralRecognition(model)
    query = "最近2年各月各个销售人员回款金额"
    sel = [{
        "id":"AUGMENTED_DATASET_LEVEL.I8a8ae5ca0178549554951b9501785cefe3f00058.MT_DIMENSION_FIELD.0c2bf3603bcd7cd295e19909c11b78fe-LEVEL-1659584084158",
        "name":"月","type":"STRING","timeLevel":"","member_amout":"12"},
        # {"id":"AUGMENTED_DATASET_LEVEL.I8a8ae5ca0178549554951b9501785cefe3f00058.MT_DIMENSION_FIELD.affebb66d15867660b4dcd4712c9a7d9-LEVEL-1667896197688",
        # "name":"季","type":"STRING","timeLevel":"month","member_amout":"1"},
        {"id":"AUGMENTED_DATASET_MEASURE.I8a8ae5ca0178549554951b9501785cefe3f00058.MT_METRICS.I8aaa80fb01804bb34bb3a22501805f86f82b02c9",
        "name":"回款金额","type":"DOUBLE","timeLevel":"","member_amout":"1"}]
    res = literal_reg.recognize(query, sel)
    print(res['origin'])
    print(res['processed'])
